{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NeuromatchAcademy/course-content/blob/master/tutorials/W2D1_DeepLearning/W2D1_Tutorial1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neuromatch Academy: Week 2, Day 1, Tutorial 1\n",
    "# Deep Learning: Decoding Neural Responses\n",
    "\n",
    "**Content creators**: Jorge A. Menendez, Carsen Stringer\n",
    "\n",
    "**Content reviewers**: Roozbeh Farhoodi,  Madineh Sarvestani, Kshitij Dwivedi, Spiros Chavlis, Ella Batty, Michael Waskom\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Tutorial Objectives\n",
    "In this tutorial, we'll use deep learning to decode stimulus information from the responses of sensory neurons. Specifically, we'll look at the activity of ~20,000 neurons in mouse primary visual cortex responding to oriented gratings recorded in [this study](https://www.biorxiv.org/content/10.1101/679324v2.abstract). Our task will be to decode the orientation of the presented stimulus from the responses of the whole population of neurons. We could do this in a number of ways, but here we'll use deep learning. Deep learning is particularly well-suited to this problem for a number of reasons:\n",
    "* The data are very high-dimensional: the neural response to a stimulus is a ~20,000 dimensional vector. Many machine learning techniques fail in such high dimensions, but deep learning actually thrives in this regime, as long as you have enough data (which we do here!).\n",
    "* As you'll be able to see below, different neurons can respond quite differently to stimuli. This complex pattern of responses will, therefore, require non-linear methods to be decoded, which we can easily do with non-linear activation functions in deep networks.\n",
    "* Deep learning architectures are highly flexible, meaning we can easily adapt the architecture of our decoding model to optimize decoding. Here, we'll focus on a single architecture, but you'll see that it can easily be modified with few changes to the code.\n",
    "\n",
    "More concretely, our goal will be learn how to:\n",
    "* Build a deep feed-forward network using PyTorch\n",
    "* Evaluate the network's outputs using PyTorch built-in loss functions\n",
    "* Compute gradients of the loss with respect to each parameter of the network using automatic differentiation\n",
    "* Implement gradient descent to optimize the network's parameters\n",
    "* (Bonus) Learn how to regularize the optimization\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:09.552092Z",
     "iopub.status.busy": "2021-05-31T20:48:09.551537Z",
     "iopub.status.idle": "2021-05-31T20:48:09.602230Z",
     "shell.execute_reply": "2021-05-31T20:48:09.601694Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 1: Decoding from neural data using feed-forward networks in pytorch\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"SlrbMvvBOzM\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Setup\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:09.606311Z",
     "iopub.status.busy": "2021-05-31T20:48:09.605791Z",
     "iopub.status.idle": "2021-05-31T20:48:10.589157Z",
     "shell.execute_reply": "2021-05-31T20:48:10.588049Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "\n",
    "import matplotlib as mpl\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:10.595825Z",
     "iopub.status.busy": "2021-05-31T20:48:10.595211Z",
     "iopub.status.idle": "2021-05-31T20:48:12.113149Z",
     "shell.execute_reply": "2021-05-31T20:48:12.111787Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Data retrieval and loading\n",
    "import hashlib\n",
    "import requests\n",
    "\n",
    "fname = \"W3D4_stringer_oribinned1.npz\"\n",
    "url = \"https://osf.io/683xc/download\"\n",
    "expected_md5 = \"436599dfd8ebe6019f066c38aed20580\"\n",
    "\n",
    "if not os.path.isfile(fname):\n",
    "  try:\n",
    "    r = requests.get(url)\n",
    "  except requests.ConnectionError:\n",
    "    print(\"!!! Failed to download data !!!\")\n",
    "  else:\n",
    "    if r.status_code != requests.codes.ok:\n",
    "      print(\"!!! Failed to download data !!!\")\n",
    "    elif hashlib.md5(r.content).hexdigest() != expected_md5:\n",
    "      print(\"!!! Data download appears corrupted !!!\")\n",
    "    else:\n",
    "      with open(fname, \"wb\") as fid:\n",
    "        fid.write(r.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:12.123044Z",
     "iopub.status.busy": "2021-05-31T20:48:12.122514Z",
     "iopub.status.idle": "2021-05-31T20:48:12.223814Z",
     "shell.execute_reply": "2021-05-31T20:48:12.223180Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Figure Settings\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:12.235703Z",
     "iopub.status.busy": "2021-05-31T20:48:12.231686Z",
     "iopub.status.idle": "2021-05-31T20:48:12.237703Z",
     "shell.execute_reply": "2021-05-31T20:48:12.237275Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Plotting Functions\n",
    "\n",
    "def plot_data_matrix(X, ax):\n",
    "  \"\"\"Visualize data matrix of neural responses using a heatmap\n",
    "\n",
    "  Args:\n",
    "    X (torch.Tensor or np.ndarray): matrix of neural responses to visualize\n",
    "        with a heatmap\n",
    "    ax (matplotlib axes): where to plot\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  cax = ax.imshow(X, cmap=mpl.cm.pink, vmin=np.percentile(X, 1), vmax=np.percentile(X, 99))\n",
    "  cbar = plt.colorbar(cax, ax=ax, label='normalized neural response')\n",
    "\n",
    "  ax.set_aspect('auto')\n",
    "  ax.set_xticks([])\n",
    "  ax.set_yticks([])\n",
    "\n",
    "def plot_decoded_results(train_loss, test_loss, test_labels, predicted_test_labels):\n",
    "  \"\"\" Plot decoding results in the form of network training loss and test predictions\n",
    "\n",
    "  Args:\n",
    "    train_loss (list): training error over iterations\n",
    "    test_labels (torch.Tensor): n_test x 1 tensor with orientations of the\n",
    "      stimuli corresponding to each row of train_data, in radians\n",
    "    predicted_test_labels (torch.Tensor): n_test x 1 tensor with predicted orientations of the\n",
    "      stimuli from decoding neural network\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  # Plot results\n",
    "  fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))\n",
    "\n",
    "  # Plot the training loss over iterations of GD\n",
    "  ax1.plot(train_loss)\n",
    "  # Plot the testing loss over iterations of GD\n",
    "  ax1.plot(test_loss)\n",
    "  ax1.legend(['train loss', 'test loss'])\n",
    "\n",
    "  # Plot true stimulus orientation vs. predicted class\n",
    "  ax2.plot(stimuli_test.squeeze(), predicted_test_labels, '.')\n",
    "\n",
    "  ax1.set_xlim([0, None])\n",
    "  ax1.set_ylim([0, None])\n",
    "  ax1.set_xlabel('iterations of gradient descent')\n",
    "  ax1.set_ylabel('negative log likelihood')\n",
    "  ax2.set_xlabel('true stimulus orientation ($^o$)')\n",
    "  ax2.set_ylabel('decoded orientation bin')\n",
    "  ax2.set_xticks(np.linspace(0, 360, n_classes + 1))\n",
    "  ax2.set_yticks(np.arange(n_classes))\n",
    "  class_bins = [f'{i * 360 / n_classes: .0f}$^o$ - {(i + 1) * 360 / n_classes: .0f}$^o$' for i in range(n_classes)]\n",
    "  ax2.set_yticklabels(class_bins);\n",
    "\n",
    "  # Draw bin edges as vertical lines\n",
    "  ax2.set_ylim(ax2.get_ylim())  # fix y-axis limits\n",
    "  for i in range(n_classes):\n",
    "    lower = i * 360 / n_classes\n",
    "    upper = (i + 1) * 360 / n_classes\n",
    "    ax2.plot([lower, lower], ax2.get_ylim(), '-', color=\"0.7\", linewidth=1, zorder=-1)\n",
    "    ax2.plot([upper, upper], ax2.get_ylim(), '-', color=\"0.7\", linewidth=1, zorder=-1)\n",
    "\n",
    "  plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:12.248610Z",
     "iopub.status.busy": "2021-05-31T20:48:12.245241Z",
     "iopub.status.idle": "2021-05-31T20:48:12.250728Z",
     "shell.execute_reply": "2021-05-31T20:48:12.250266Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Helper Functions\n",
    "\n",
    "def load_data(data_name=fname, bin_width=1):\n",
    "  \"\"\"Load mouse V1 data from Stringer et al. (2019)\n",
    "\n",
    "  Data from study reported in this preprint:\n",
    "  https://www.biorxiv.org/content/10.1101/679324v2.abstract\n",
    "\n",
    "  These data comprise time-averaged responses of ~20,000 neurons\n",
    "  to ~4,000 stimulus gratings of different orientations, recorded\n",
    "  through Calcium imaging. The responses have been normalized by\n",
    "  spontaneous levels of activity and then z-scored over stimuli, so\n",
    "  expect negative numbers. They have also been binned and averaged\n",
    "  to each degree of orientation.\n",
    "\n",
    "  This function returns the relevant data (neural responses and\n",
    "  stimulus orientations) in a torch.Tensor of data type torch.float32\n",
    "  in order to match the default data type for nn.Parameters in\n",
    "  Google Colab.\n",
    "\n",
    "  This function will actually average responses to stimuli with orientations\n",
    "  falling within bins specified by the bin_width argument. This helps\n",
    "  produce individual neural \"responses\" with smoother and more\n",
    "  interpretable tuning curves.\n",
    "\n",
    "  Args:\n",
    "    bin_width (float): size of stimulus bins over which to average neural\n",
    "      responses\n",
    "\n",
    "  Returns:\n",
    "    resp (torch.Tensor): n_stimuli x n_neurons matrix of neural responses,\n",
    "        each row contains the responses of each neuron to a given stimulus.\n",
    "        As mentioned above, neural \"response\" is actually an average over\n",
    "        responses to stimuli with similar angles falling within specified bins.\n",
    "    stimuli: (torch.Tensor): n_stimuli x 1 column vector with orientation\n",
    "        of each stimulus, in degrees. This is actually the mean orientation\n",
    "        of all stimuli in each bin.\n",
    "\n",
    "  \"\"\"\n",
    "  with np.load(data_name) as dobj:\n",
    "    data = dict(**dobj)\n",
    "  resp = data['resp']\n",
    "  stimuli = data['stimuli']\n",
    "\n",
    "  if bin_width > 1:\n",
    "    # Bin neural responses and stimuli\n",
    "    bins = np.digitize(stimuli, np.arange(0, 360 + bin_width, bin_width))\n",
    "    stimuli_binned = np.array([stimuli[bins == i].mean() for i in np.unique(bins)])\n",
    "    resp_binned = np.array([resp[bins == i, :].mean(0) for i in np.unique(bins)])\n",
    "  else:\n",
    "    resp_binned = resp\n",
    "    stimuli_binned = stimuli\n",
    "\n",
    "  # Return as torch.Tensor\n",
    "  resp_tensor = torch.tensor(resp_binned, dtype=torch.float32)\n",
    "  stimuli_tensor = torch.tensor(stimuli_binned, dtype=torch.float32).unsqueeze(1)  # add singleton dimension to make a column vector\n",
    "\n",
    "  return resp_tensor, stimuli_tensor\n",
    "\n",
    "\n",
    "def identityLine():\n",
    "  \"\"\"\n",
    "  Plot the identity line y=x\n",
    "  \"\"\"\n",
    "  ax = plt.gca()\n",
    "  lims = np.array([ax.get_xlim(), ax.get_ylim()])\n",
    "  minval = lims[:, 0].min()\n",
    "  maxval = lims[:, 1].max()\n",
    "  equal_lims = [minval, maxval]\n",
    "  ax.set_xlim(equal_lims)\n",
    "  ax.set_ylim(equal_lims)\n",
    "  line = ax.plot([minval, maxval], [minval, maxval], color=\"0.7\")\n",
    "  line[0].set_zorder(-1)\n",
    "\n",
    "def get_data(n_stim, train_data, train_labels):\n",
    "  \"\"\" Return n_stim randomly drawn stimuli/resp pairs\n",
    "\n",
    "  Args:\n",
    "    n_stim (scalar): number of stimuli to draw\n",
    "    resp (torch.Tensor):\n",
    "    train_data (torch.Tensor): n_train x n_neurons tensor with neural\n",
    "      responses to train on\n",
    "    train_labels (torch.Tensor): n_train x 1 tensor with orientations of the\n",
    "      stimuli corresponding to each row of train_data, in radians\n",
    "\n",
    "  Returns:\n",
    "    (torch.Tensor, torch.Tensor): n_stim x n_neurons tensor of neural responses and n_stim x 1 of orientations respectively\n",
    "  \"\"\"\n",
    "  n_stimuli = train_labels.shape[0]\n",
    "  istim = np.random.choice(n_stimuli, n_stim)\n",
    "  r = train_data[istim]  # neural responses to this stimulus\n",
    "  ori = train_labels[istim]  # true stimulus orientation\n",
    "\n",
    "  return r, ori\n",
    "\n",
    "def stimulus_class(ori, n_classes):\n",
    "  \"\"\"Get stimulus class from stimulus orientation\n",
    "\n",
    "  Args:\n",
    "    ori (torch.Tensor): orientations of stimuli to return classes for\n",
    "    n_classes (int): total number of classes\n",
    "\n",
    "  Returns:\n",
    "    torch.Tensor: 1D tensor with the classes for each stimulus\n",
    "\n",
    "  \"\"\"\n",
    "  bins = np.linspace(0, 360, n_classes + 1)\n",
    "  return torch.tensor(np.digitize(ori.squeeze(), bins)) - 1  # minus 1 to accomodate Python indexing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 1: Load and visualize data\n",
    "\n",
    "In the next cell, we have provided code to load the data and plot the matrix of neural responses.\n",
    "\n",
    "Next to it, we plot the tuning curves of three randomly selected neurons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:12.256054Z",
     "iopub.status.busy": "2021-05-31T20:48:12.255486Z",
     "iopub.status.idle": "2021-05-31T20:48:14.698512Z",
     "shell.execute_reply": "2021-05-31T20:48:14.698987Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "#@markdown Execute this cell to load and visualize data\n",
    "\n",
    "# Load data\n",
    "resp_all, stimuli_all = load_data()  # argument to this function specifies bin width\n",
    "n_stimuli, n_neurons = resp_all.shape\n",
    "\n",
    "print(f'{n_neurons} neurons in response to {n_stimuli} stimuli')\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(2 * 6, 5))\n",
    "\n",
    "# Visualize data matrix\n",
    "plot_data_matrix(resp_all[:100, :].T, ax1)  # plot responses of first 100 neurons\n",
    "ax1.set_xlabel('stimulus')\n",
    "ax1.set_ylabel('neuron')\n",
    "\n",
    "# Plot tuning curves of three random neurons\n",
    "ineurons = np.random.choice(n_neurons, 3, replace=False)  # pick three random neurons\n",
    "ax2.plot(stimuli_all, resp_all[:, ineurons])\n",
    "ax2.set_xlabel('stimulus orientation ($^o$)')\n",
    "ax2.set_ylabel('neural response')\n",
    "ax2.set_xticks(np.linspace(0, 360, 5))\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will split our data into a training set and test set. In particular, we will have a training set of orientations (`stimuli_train`) and the corresponding responses (`resp_train`). Our testing set will have held-out orientations (`stimuli_test`) and the corresponding responses (`resp_test`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:14.704061Z",
     "iopub.status.busy": "2021-05-31T20:48:14.703512Z",
     "iopub.status.idle": "2021-05-31T20:48:14.716253Z",
     "shell.execute_reply": "2021-05-31T20:48:14.717018Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "#@markdown Execute this cell to split into training and test sets\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(4)\n",
    "torch.manual_seed(4)\n",
    "\n",
    "# Split data into training set and testing set\n",
    "n_train = int(0.6 * n_stimuli)  # use 60% of all data for training set\n",
    "ishuffle = torch.randperm(n_stimuli)\n",
    "itrain = ishuffle[:n_train]  # indices of data samples to include in training set\n",
    "itest = ishuffle[n_train:]  # indices of data samples to include in testing set\n",
    "stimuli_test = stimuli_all[itest]\n",
    "resp_test = resp_all[itest]\n",
    "stimuli_train = stimuli_all[itrain]\n",
    "resp_train = resp_all[itrain]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 2: Deep feed-forward networks in *pytorch* \n",
    "\n",
    "We'll now build a simple deep neural network that takes as input a vector of neural responses and outputs a single number representing the decoded stimulus orientation.\n",
    "\n",
    "To keep things simple, we'll build a deep network with **one** hidden layer that is **fully connected**, which means *every* output is a weighted sum of *every* input. See Bonus Section 1 for a deeper discussion of what this choice entails, and when one might want to use deeper/shallower and wider/narrower architectures.\n",
    "\n",
    "Let $\\mathbf{r}^{(n)} = \\begin{bmatrix} r_1^{(n)} & r_2^{(n)} & \\ldots & r_N^{(n)} \\end{bmatrix}^T$ denote the vector of neural responses (of neurons $1, \\ldots, N$) to the $n$th stimulus. The network we will use is described by the following set of equations:\n",
    "\\begin{align}\n",
    "    \\mathbf{h}^{(n)} &= \\mathbf{W}^{in} \\mathbf{r}^{(n)} + \\mathbf{b}^{in}, && [\\mathbf{W}^{in}: M \\times N], \\\\\n",
    "    y^{(n)} &= \\mathbf{W}^{out} \\mathbf{h}^{(n)} + \\mathbf{b}^{out},  && [\\mathbf{W}^{out}: 1 \\times M],\n",
    "\\end{align}\n",
    "where $y^{(n)}$ denotes the scalar output of the network: the decoded orientation of the $n$th stimulus. \n",
    "\n",
    "The $M$-dimensional vector $\\mathbf{h}^{(n)}$ denotes the activations of the **hidden layer** of the network. The blue components of this diagram denote the **parameters** of the network, which we will later optimize with gradient descent. These include all the weights and biases $\\mathbf{W}^{in}, \\mathbf{b}^{in}, \\mathbf{W}^{out}, \\mathbf{b}^{out}$.\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"https://github.com/NeuromatchAcademy/course-content/blob/master/tutorials/static/one-layer-network.png?raw=true\" width=\"450\" />\n",
    "</p>\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 2.1: Introduction to PyTorch\n",
    "\n",
    "Here, we'll use the **PyTorch** package to build, run, and train deep networks of this form in Python. There are two core components to the PyTorch package: \n",
    "\n",
    "1. The first is the `torch.Tensor` data type used in PyTorch. `torch.Tensor`'s are effectively just like a `numpy` arrays, except that they have some important attributes and methods needed for automatic differentiation (to be discussed below). They also come along with infrastructure for easily storing and computing with them on GPU's, a capability we won't touch on here but which can be really useful in practice.\n",
    "\n",
    "2. The second core ingredient is the PyTorch `nn.Module` class. This is the class we'll use for constructing deep networks, so that we can then easily train them using built-in PyTorch functions. Keep in my mind that `nn.Module` classes can actually be used to build, run, and train any model -- not just deep networks!\n",
    "\n",
    "  The next cell contains code for building the deep network we defined above using the `nn.Module` class. It contains three key ingredients:\n",
    "\n",
    "  * `__init__()` method to initialize its parameters, like in any other Python class. In this case, it takes two arguments:\n",
    "    * `n_inputs`: the number of input units. This should always be set to the number of neurons whose activities are being decoded (i.e. the dimensionality of the input to the network). \n",
    "    * `n_hidden`: the number of hidden units. This is a parameter that we are free to vary in deciding how to build our network. See Bonus Section 1 for a discussion of how this architectural choice affects the computations the network can perform.\n",
    "\n",
    "  * `nn.Linear` modules, which are built-in PyTorch classes containing all the weights and biases for a given network layer (documentation [here](https://pytorch.org/docs/master/generated/torch.nn.Linear.html)). This class takes two arguments to initialize:\n",
    "    * \\# of inputs to that layer\n",
    "    * \\# of outputs from that layer\n",
    "\n",
    "    For the input layer, for example, we have:\n",
    "    * \\# of inputs = \\# of neurons whose responses are to be decoded ($N$, specified by `n_inputs`)\n",
    "    * \\# of outputs = \\# of hidden layer units ($M$, specified by `n_hidden`)\n",
    "    \n",
    "    PyTorch will initialize all weights and biases randomly.\n",
    "\n",
    "  * `forward()` method, which takes as argument an input to the network and returns the network output. In our case, this comprises computing the output $y$ from a given input $\\mathbf{r}$ using the above two equations. See the next cell for code implementing this computation using the built-in PyTorch `nn.Linear` classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:14.722985Z",
     "iopub.status.busy": "2021-05-31T20:48:14.722446Z",
     "iopub.status.idle": "2021-05-31T20:48:14.725881Z",
     "shell.execute_reply": "2021-05-31T20:48:14.725410Z"
    }
   },
   "outputs": [],
   "source": [
    "class DeepNet(nn.Module):\n",
    "  \"\"\"Deep Network with one hidden layer\n",
    "\n",
    "  Args:\n",
    "    n_inputs (int): number of input units\n",
    "    n_hidden (int): number of units in hidden layer\n",
    "\n",
    "  Attributes:\n",
    "    in_layer (nn.Linear): weights and biases of input layer\n",
    "    out_layer (nn.Linear): weights and biases of output layer\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  def __init__(self, n_inputs, n_hidden):\n",
    "    super().__init__()  # needed to invoke the properties of the parent class nn.Module\n",
    "    self.in_layer = nn.Linear(n_inputs, n_hidden) # neural activity --> hidden units\n",
    "    self.out_layer = nn.Linear(n_hidden, 1) # hidden units --> output\n",
    "\n",
    "  def forward(self, r):\n",
    "    \"\"\"Decode stimulus orientation from neural responses\n",
    "\n",
    "    Args:\n",
    "      r (torch.Tensor): vector of neural responses to decode, must be of\n",
    "        length n_inputs. Can also be a tensor of shape n_stimuli x n_inputs,\n",
    "        containing n_stimuli vectors of neural responses\n",
    "\n",
    "    Returns:\n",
    "      torch.Tensor: network outputs for each input provided in r. If\n",
    "        r is a vector, then y is a 1D tensor of length 1. If r is a 2D\n",
    "        tensor then y is a 2D tensor of shape n_stimuli x 1.\n",
    "\n",
    "    \"\"\"\n",
    "    h = self.in_layer(r)  # hidden representation\n",
    "    y = self.out_layer(h)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell contains code for initializing and running this network. We use it to decode stimulus orientation from a vector of neural responses to the very first stimulus. Note that when the initialized network class is called as a function on an input (e.g. `net(r)`), its `.forward()` method is called. This is a special property of the `nn.Module` class.\n",
    "\n",
    "Note that the decoded orientations at this point will be nonsense, since the network has been initialized with random weights. Below, we'll learn how to optimize these weights for good stimulus decoding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:14.730863Z",
     "iopub.status.busy": "2021-05-31T20:48:14.730335Z",
     "iopub.status.idle": "2021-05-31T20:48:14.773812Z",
     "shell.execute_reply": "2021-05-31T20:48:14.773322Z"
    }
   },
   "outputs": [],
   "source": [
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "# Initialize a deep network with M=200 hidden units\n",
    "net = DeepNet(n_neurons, 200)\n",
    "\n",
    "# Get neural responses (r) to and orientation (ori) to one stimulus in dataset\n",
    "r, ori = get_data(1, resp_train, stimuli_train)  # using helper function get_data\n",
    "\n",
    "# Decode orientation from these neural responses using initialized network\n",
    "out = net(r)  # compute output from network, equivalent to net.forward(r)\n",
    "\n",
    "print('decoded orientation: %.2f degrees' % out)\n",
    "print('true orientation: %.2f degrees' % ori)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "### Section 2.2: Activation functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:14.779090Z",
     "iopub.status.busy": "2021-05-31T20:48:14.778498Z",
     "iopub.status.idle": "2021-05-31T20:48:14.791772Z",
     "shell.execute_reply": "2021-05-31T20:48:14.792205Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 2: Nonlinear activation functions\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"JAdukDCQALA\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the deep network we constructed above comprises solely **linear** operations on each layer: each layer is just a weighted sum of all the elements in the previous layer. It turns out that linear hidden layers like this aren't particularly useful, since a sequence of linear transformations is actually essentially the same as a single linear transformation. We can see this from the above equations by plugging in the first one into the second one to obtain\n",
    "\\begin{equation}\n",
    "    y^{(n)} = \\mathbf{W}^{out} \\left( \\mathbf{W}^{in} \\mathbf{r}^{(n)} + \\mathbf{b}^{in} \\right) + \\mathbf{b}^{out} = \\mathbf{W}^{out}\\mathbf{W}^{in} \\mathbf{r}^{(n)} + \\left( \\mathbf{W}^{out}\\mathbf{b}^{in} + \\mathbf{b}^{out} \\right)\n",
    "\\end{equation}\n",
    "In other words, the output is still just a weighted sum of elements in the input -- the hidden layer has done nothing to change this.\n",
    "\n",
    "To extend the set of computable input/output transformations to more than just weighted sums, we'll incorporate a **non-linear activation function** in the hidden units. This is done by simply modifying the equation for the hidden layer activations to be\n",
    "\\begin{equation}\n",
    "    \\mathbf{h}^{(n)} = \\phi(\\mathbf{W}^{in} \\mathbf{r}^{(n)} + \\mathbf{b}^{in})\n",
    "\\end{equation}\n",
    "where $\\phi$ is referred to as the activation function. Using a non-linear activation function will ensure that the hidden layer performs a non-linear transformation of the input, which will make our network much more powerful (or *expressive*, see Bonus Section 1). In practice, deep networks *always* use non-linear activation functions.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Coding Exercise 2.2: Nonlinear Activations \n",
    "\n",
    "Create a new class `DeepNetReLU` by modifying our above deep network model to use a non-linear activation function. We'll use the linear rectification function:\n",
    "\\begin{equation}\n",
    "  \\phi(x) = \n",
    "  \\begin{cases}\n",
    "    x & \\text{if } x > 0 \\\\\n",
    "    0 & \\text{else}\n",
    "  \\end{cases}\n",
    "\\end{equation}\n",
    "which can be implemented in PyTorch using `torch.relu()`. Hidden layers with this activation function are typically referred to as \"**Re**ctified **L**inear **U**nits\", or **ReLU**'s.\n",
    "\n",
    "Initialize this network with 10 hidden units and run on an example stimulus.\n",
    "\n",
    "**Hint**: you only need to modify the `forward()` method of the above `DeepNet()` class.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:14.798538Z",
     "iopub.status.busy": "2021-05-31T20:48:14.797983Z",
     "iopub.status.idle": "2021-05-31T20:48:14.911647Z",
     "shell.execute_reply": "2021-05-31T20:48:14.909539Z"
    }
   },
   "outputs": [],
   "source": [
    "class DeepNetReLU(nn.Module):\n",
    "  \"\"\" network with a single hidden layer h with a RELU \"\"\"\n",
    "\n",
    "  def __init__(self, n_inputs, n_hidden):\n",
    "    super().__init__()  # needed to invoke the properties of the parent class nn.Module\n",
    "    self.in_layer = nn.Linear(n_inputs, n_hidden) # neural activity --> hidden units\n",
    "    self.out_layer = nn.Linear(n_hidden, 1) # hidden units --> output\n",
    "\n",
    "  def forward(self, r):\n",
    "\n",
    "    ############################################################################\n",
    "    ## TO DO for students: write code for computing network output using a\n",
    "    ## rectified linear activation function for the hidden units\n",
    "    # Fill out function and remove\n",
    "    raise NotImplementedError(\"Student exercise: complete DeepNetReLU forward\")\n",
    "    ############################################################################\n",
    "\n",
    "    h = ...\n",
    "    y = ...\n",
    "\n",
    "    return y\n",
    "\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "# Get neural responses (r) to and orientation (ori) to one stimulus in dataset\n",
    "r, ori = get_data(1, resp_train, stimuli_train)\n",
    "\n",
    "# Initialize deep network with M=10 hidden units and uncomment lines below\n",
    "net = DeepNetReLU(n_neurons, 10)\n",
    "\n",
    "# Decode orientation from these neural responses using initialized network\n",
    "# net(r) is equivalent to net.forward(r)\n",
    "out = net(r)\n",
    "\n",
    "print('decoded orientation: %.2f degrees' % out)\n",
    "print('true orientation: %.2f degrees' % ori)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:14.918589Z",
     "iopub.status.busy": "2021-05-31T20:48:14.917460Z",
     "iopub.status.idle": "2021-05-31T20:48:14.927037Z",
     "shell.execute_reply": "2021-05-31T20:48:14.927546Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "\n",
    "class DeepNetReLU(nn.Module):\n",
    "  \"\"\" network with a single hidden layer h with a RELU \"\"\"\n",
    "\n",
    "  def __init__(self, n_inputs, n_hidden):\n",
    "    super().__init__()  # needed to invoke the properties of the parent class nn.Module\n",
    "    self.in_layer = nn.Linear(n_inputs, n_hidden) # neural activity --> hidden units\n",
    "    self.out_layer = nn.Linear(n_hidden, 1) # hidden units --> output\n",
    "\n",
    "  def forward(self, r):\n",
    "\n",
    "    h = torch.relu(self.in_layer(r))\n",
    "    y = self.out_layer(h)\n",
    "\n",
    "    return y\n",
    "\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "# Get neural responses (r) to and orientation (ori) to one stimulus in dataset\n",
    "r, ori = get_data(1, resp_train, stimuli_train)\n",
    "\n",
    "# Initialize deep network with M=10 hidden units and uncomment lines below\n",
    "net = DeepNetReLU(n_neurons, 10)\n",
    "\n",
    "# Decode orientation from these neural responses using initialized network\n",
    "# net(r) is equivalent to net.forward(r)\n",
    "out = net(r)\n",
    "\n",
    "print('decoded orientation: %.2f degrees' % out)\n",
    "print('true orientation: %.2f degrees' % ori)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see that the decoded orientation is 0.13 $^{\\circ}$ while the true orientation is 139.00 $^{\\circ}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 3: Loss functions and gradient descent\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:14.933233Z",
     "iopub.status.busy": "2021-05-31T20:48:14.932344Z",
     "iopub.status.idle": "2021-05-31T20:48:14.947027Z",
     "shell.execute_reply": "2021-05-31T20:48:14.946543Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 3: Loss functions & gradient descent\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"aEtKpzEuviw\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 3.1: Loss functions\n",
    "\n",
    "Because the weights of the network are currently randomly chosen, the outputs of the network are nonsense: the decoded stimulus orientation is nowhere close to the true stimulus orientation. We'll shortly write some code to change these weights so that the network does a better job of decoding.\n",
    "\n",
    "But to do so, we first need to define what we mean by \"better\". One simple way of defining this is to use the squared error\n",
    "\\begin{equation}\n",
    "    L = (y - \\tilde{y})^2\n",
    "\\end{equation}\n",
    "where $y$ is the network output and $\\tilde{y}$ is the true stimulus orientation. When the decoded stimulus orientation is far from the true stimulus orientation, $L$ will be large. We thus refer to $L$ as the **loss function**, as it quantifies how *bad* the network is at decoding stimulus orientation.\n",
    "\n",
    "PyTorch actually carries with it a number of built-in loss functions. The one corresponding to the squared error is called `nn.MSELoss()`. This will take as arguments a **batch** of network outputs $y_1, y_2, \\ldots, y_P$ and corresponding target outputs $\\tilde{y}_1, \\tilde{y}_2, \\ldots, \\tilde{y}_P$, and compute the **mean squared error (MSE)**\n",
    "\\begin{equation}\n",
    "    L = \\frac{1}{P}\\sum_{n=1}^P \\left(y^{(n)} - \\tilde{y}^{(n)}\\right)^2\n",
    "\\end{equation}\n",
    "where $P$ is the number of different stimuli in a batch, called the *batch size*.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Coding Exercise 3.1: Computing MSE \n",
    "\n",
    "\n",
    "Evaluate the mean squared error for a deep network with $M=10$ rectified linear units, on the decoded orientations from neural responses to 20 random stimuli."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:14.951705Z",
     "iopub.status.busy": "2021-05-31T20:48:14.951133Z",
     "iopub.status.idle": "2021-05-31T20:48:14.959247Z",
     "shell.execute_reply": "2021-05-31T20:48:14.958760Z"
    }
   },
   "outputs": [],
   "source": [
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "# Initialize a deep network with M=10 hidden units\n",
    "net = DeepNetReLU(n_neurons, 10)\n",
    "\n",
    "# Get neural responses to first 20 stimuli in the data set\n",
    "r, ori = get_data(20, resp_train, stimuli_train)\n",
    "\n",
    "# Decode orientation from these neural responses\n",
    "out = net(r)\n",
    "\n",
    "###################################################\n",
    "## TO DO for students: evaluate mean squared error\n",
    "###################################################\n",
    "\n",
    "# Initialize PyTorch mean squared error loss function (Hint: look at nn.MSELoss)\n",
    "loss_fn = ...\n",
    "\n",
    "# Evaluate mean squared error\n",
    "loss = ...\n",
    "\n",
    "# Uncomment once above is filled in\n",
    "# print('mean squared error: %.2f' % loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:14.963790Z",
     "iopub.status.busy": "2021-05-31T20:48:14.963264Z",
     "iopub.status.idle": "2021-05-31T20:48:14.971423Z",
     "shell.execute_reply": "2021-05-31T20:48:14.970921Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "# Initialize a deep network with M=10 hidden units\n",
    "net = DeepNetReLU(n_neurons, 10)\n",
    "\n",
    "# Get neural responses to first 20 stimuli in the data set\n",
    "r, ori = get_data(20, resp_train, stimuli_train)\n",
    "\n",
    "# Decode orientation from these neural responses\n",
    "out = net(r)\n",
    "\n",
    "# Initialize PyTorch mean squared error loss function (Hint: look at nn.MSELoss)\n",
    "loss_fn = nn.MSELoss()\n",
    "\n",
    "# Evaluate mean squared error\n",
    "loss = loss_fn(out, ori)\n",
    "\n",
    "print('mean squared error: %.2f' % loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see a mean squared error of 42949.14."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "### Section 3.2: Optimization with gradient descent\n",
    "\n",
    "Our goal is now to modify the weights to make the mean squared error loss $L$ as small as possible over the whole data set. To do this, we'll use the **gradient descent (GD)** algorithm, which consists of iterating three steps:\n",
    "1. **Evaluate the loss** on the training data,\n",
    "```\n",
    "out = net(train_data)\n",
    "loss = loss_fn(out, train_labels)\n",
    "```\n",
    "where `train_data` are the network inputs in the training data (in our case, neural responses), and `train_labels` are the target outputs for each input (in our case, true stimulus orientations).\n",
    "2. **Compute the gradient of the loss** with respect to each of the network weights. In PyTorch, we can do this with one line of code:\n",
    "```\n",
    "loss.backward()\n",
    "```\n",
    "3. **Update the network weights** by descending the gradient. In Pytorch, we can do this using built-in optimizers. We'll use the `optim.SGD` optimizer (documentation [here](https://pytorch.org/docs/stable/optim.html#torch.optim.SGD)) which updates parameters along the negative gradient, scaled by a learning rate (see Bonus Section 2.1 for details). To initialize this optimizer, we have to tell it\n",
    "  * which parameters to update, and\n",
    "  * what learning rate to use\n",
    "\n",
    "  For example, to optimize *all* the parameters of a network `net` using a learning rate of .001, the optimizer would be initialized as follows\n",
    "  ```\n",
    "  optimizer = optim.SGD(net.parameters(), lr=.001)\n",
    "  ```\n",
    "  where `.parameters()` is a method of the `nn.Module` class that returns a [Python generator object](https://wiki.python.org/moin/Generators) over all the parameters of that `nn.Module` class (in our case, $\\mathbf{W}^{in}, \\mathbf{b}^{in}, \\mathbf{W}^{out}, \\mathbf{b}^{out}$).\n",
    "  \n",
    "  After computing all the parameter gradients in step 2, we can then update each of these parameters using the `.step()` method of this optimizer,\n",
    "  ```\n",
    "  optimizer.step()\n",
    "  ```\n",
    "\n",
    "Finally, an important detail to remember is that the gradients of each parameter need to be cleared before calling `.backward()`, or else PyTorch will try to accumulate gradients across iterations. This can again be done using built-in optimizers via the method `zero_grad()`, as follows:\n",
    "```\n",
    "optimizer.zero_grad()\n",
    "```\n",
    "\n",
    "Putting all this together, each iteration of the GD algorith will contain a block of code that looks something like this:\n",
    "```\n",
    "Get outputs from network\n",
    "Evaluate loss\n",
    "\n",
    "# Compute gradients\n",
    "optimizer.zero_grad()  # clear gradients\n",
    "loss.backward()\n",
    "\n",
    "# Update weights\n",
    "optimizer.step()\n",
    "```\n",
    "\n",
    "In the next exercise, we'll give you a code skeleton for implementing the GD algorithm. Your job will be to fill in the blanks.\n",
    "\n",
    "For the mathematical details of the GD algorithm, see Bonus Section 2.1. Note, in particular, that here we using the gradient descent algorithm, rather than the more commonly used *stochastic* gradient descent algorithm. See Bonus Section 2.2 for a more detailed discussion of how these differ and when one might need to use the stochastic variant."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Coding Exercise 3.2: Gradient descent in PyTorch\n",
    "\n",
    "Complete the function `train()` that uses the gradient descent algorithm to optimize the weights of a given network. This function takes as input arguments\n",
    "* `net`: the PyTorch network whose weights to optimize\n",
    "* `loss_fn`: the PyTorch loss function to use to evaluate the loss\n",
    "* `train_data`: the training data to evaluate the loss on (i.e. neural responses to decode)\n",
    "* `train_labels`: the target outputs for each data point in `train_data` (i.e. true stimulus orientations)\n",
    "\n",
    "We will then train a neural network on our data and plot the loss (mean squared error) over time. When we run this function, behind the scenes PyTorch is actually changing the parameters inside this network to make the network better at decoding, so its weights will now be different than they were at initialization.\n",
    "\n",
    "\n",
    "**Hint:** all the code you need for doing this is provided in the above description of the GD algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:14.983588Z",
     "iopub.status.busy": "2021-05-31T20:48:14.982926Z",
     "iopub.status.idle": "2021-05-31T20:48:14.996489Z",
     "shell.execute_reply": "2021-05-31T20:48:14.995925Z"
    }
   },
   "outputs": [],
   "source": [
    "def train(net, loss_fn, train_data, train_labels,\n",
    "          n_iter=50, learning_rate=1e-4,\n",
    "          test_data=None, test_labels=None,\n",
    "          L2_penalty=0, L1_penalty=0):\n",
    "  \"\"\"Run gradient descent to opimize parameters of a given network\n",
    "\n",
    "  Args:\n",
    "    net (nn.Module): PyTorch network whose parameters to optimize\n",
    "    loss_fn: built-in PyTorch loss function to minimize\n",
    "    train_data (torch.Tensor): n_train x n_neurons tensor with neural\n",
    "      responses to train on\n",
    "    train_labels (torch.Tensor): n_train x 1 tensor with orientations of the\n",
    "      stimuli corresponding to each row of train_data\n",
    "    n_iter (int, optional): number of iterations of gradient descent to run\n",
    "    learning_rate (float, optional): learning rate to use for gradient descent\n",
    "    test_data (torch.Tensor, optional): n_test x n_neurons tensor with neural\n",
    "      responses to test on\n",
    "    test_labels (torch.Tensor, optional): n_test x 1 tensor with orientations of\n",
    "      the stimuli corresponding to each row of test_data\n",
    "    L2_penalty (float, optional): l2 penalty regularizer coefficient\n",
    "    L1_penalty (float, optional): l1 penalty regularizer coefficient\n",
    "\n",
    "  Returns:\n",
    "    (list): training loss over iterations\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  # Initialize PyTorch SGD optimizer\n",
    "  optimizer = optim.SGD(net.parameters(), lr=learning_rate)\n",
    "\n",
    "  # Placeholder to save the loss at each iteration\n",
    "  train_loss = []\n",
    "  test_loss = []\n",
    "\n",
    "  # Loop over epochs\n",
    "  for i in range(n_iter):\n",
    "\n",
    "    ######################################################################\n",
    "    ## TO DO for students: fill in missing code for GD iteration\n",
    "    raise NotImplementedError(\"Student exercise: write code for GD iterations\")\n",
    "    ######################################################################\n",
    "\n",
    "    # compute network output from inputs in train_data\n",
    "    out = ...  # compute network output from inputs in train_data\n",
    "\n",
    "    # evaluate loss function\n",
    "    if L2_penalty==0 and L1_penalty==0:\n",
    "      # normal loss function\n",
    "      loss = loss_fn(out, train_labels)\n",
    "    else:\n",
    "      # custom loss function from bonus exercise 3.3\n",
    "      loss = loss_fn(out, train_labels, net.in_layer.weight,\n",
    "                     L2_penalty, L1_penalty)\n",
    "\n",
    "    # Clear previous gradients\n",
    "    ...\n",
    "\n",
    "    # Compute gradients\n",
    "    ...\n",
    "\n",
    "    # Update weights\n",
    "    ...\n",
    "\n",
    "    # Store current value of loss\n",
    "    train_loss.append(loss.item())  # .item() needed to transform the tensor output of loss_fn to a scalar\n",
    "\n",
    "    # Get loss for test_data, if given (we will use this in the bonus exercise 3.2 and 3.3)\n",
    "    if test_data is not None:\n",
    "      out_test = net(test_data)\n",
    "      # evaluate loss function\n",
    "      if L2_penalty==0 and L1_penalty==0:\n",
    "        # normal loss function\n",
    "        loss_test = loss_fn(out_test, test_labels)\n",
    "      else:\n",
    "        # (BONUS code) custom loss function from Bonus exercise 3.3\n",
    "        loss_test = loss_fn(out_test, test_labels, net.in_layer.weight,\n",
    "                            L2_penalty, L1_penalty)\n",
    "      test_loss.append(loss_test.item())  # .item() needed to transform the tensor output of loss_fn to a scalar\n",
    "\n",
    "    # Track progress\n",
    "    if (i + 1) % (n_iter // 5) == 0:\n",
    "      if test_data is None:\n",
    "        print(f'iteration {i + 1}/{n_iter} | loss: {loss.item():.3f}')\n",
    "      else:\n",
    "        print(f'iteration {i + 1}/{n_iter} | loss: {loss.item():.3f} | test_loss: {loss_test.item():.3f}')\n",
    "\n",
    "  if test_data is None:\n",
    "    return train_loss\n",
    "  else:\n",
    "    return train_loss, test_loss\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "# Initialize network with 10 hidden units\n",
    "net = DeepNetReLU(n_neurons, 10)\n",
    "\n",
    "# Initialize built-in PyTorch MSE loss function\n",
    "loss_fn = nn.MSELoss()\n",
    "\n",
    "# Run GD on data\n",
    "train_loss = train(net, loss_fn, resp_train, stimuli_train)\n",
    "\n",
    "# Plot the training loss over iterations of GD\n",
    "with plt.xkcd():\n",
    "  plt.plot(train_loss)\n",
    "  plt.xlim([0, None])\n",
    "  plt.ylim([0, None])\n",
    "  plt.xlabel('iterations of gradient descent')\n",
    "  plt.ylabel('mean squared error')\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:15.008453Z",
     "iopub.status.busy": "2021-05-31T20:48:15.007843Z",
     "iopub.status.idle": "2021-05-31T20:48:15.441051Z",
     "shell.execute_reply": "2021-05-31T20:48:15.440572Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "\n",
    "def train(net, loss_fn, train_data, train_labels,\n",
    "          n_iter=50, learning_rate=1e-4,\n",
    "          test_data=None, test_labels=None,\n",
    "          L2_penalty=0, L1_penalty=0):\n",
    "  \"\"\"Run gradient descent to opimize parameters of a given network\n",
    "\n",
    "  Args:\n",
    "    net (nn.Module): PyTorch network whose parameters to optimize\n",
    "    loss_fn: built-in PyTorch loss function to minimize\n",
    "    train_data (torch.Tensor): n_train x n_neurons tensor with neural\n",
    "      responses to train on\n",
    "    train_labels (torch.Tensor): n_train x 1 tensor with orientations of the\n",
    "      stimuli corresponding to each row of train_data\n",
    "    n_iter (int, optional): number of iterations of gradient descent to run\n",
    "    learning_rate (float, optional): learning rate to use for gradient descent\n",
    "    test_data (torch.Tensor, optional): n_test x n_neurons tensor with neural\n",
    "      responses to test on\n",
    "    test_labels (torch.Tensor, optional): n_test x 1 tensor with orientations of\n",
    "      the stimuli corresponding to each row of test_data\n",
    "    L2_penalty (float, optional): l2 penalty regularizer coefficient\n",
    "    L1_penalty (float, optional): l1 penalty regularizer coefficient\n",
    "\n",
    "  Returns:\n",
    "    (list): training loss over iterations\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  # Initialize PyTorch SGD optimizer\n",
    "  optimizer = optim.SGD(net.parameters(), lr=learning_rate)\n",
    "\n",
    "  # Placeholder to save the loss at each iteration\n",
    "  train_loss = []\n",
    "  test_loss = []\n",
    "\n",
    "  # Loop over epochs\n",
    "  for i in range(n_iter):\n",
    "\n",
    "    # compute network output from inputs in train_data\n",
    "    out = net(train_data)  # compute network output from inputs in train_data\n",
    "\n",
    "    # evaluate loss function\n",
    "    if L2_penalty==0 and L1_penalty==0:\n",
    "      # normal loss function\n",
    "      loss = loss_fn(out, train_labels)\n",
    "    else:\n",
    "      # custom loss function from bonus exercise 3.3\n",
    "      loss = loss_fn(out, train_labels, net.in_layer.weight,\n",
    "                     L2_penalty, L1_penalty)\n",
    "\n",
    "    # Clear previous gradients\n",
    "    optimizer.zero_grad()\n",
    "    # Compute gradients\n",
    "    loss.backward()\n",
    "\n",
    "    # Update weights\n",
    "    optimizer.step()\n",
    "\n",
    "    # Store current value of loss\n",
    "    train_loss.append(loss.item())  # .item() needed to transform the tensor output of loss_fn to a scalar\n",
    "\n",
    "    # Get loss for test_data, if given (we will use this in the bonus exercise 3.2 and 3.3)\n",
    "    if test_data is not None:\n",
    "      out_test = net(test_data)\n",
    "      # evaluate loss function\n",
    "      if L2_penalty==0 and L1_penalty==0:\n",
    "        # normal loss function\n",
    "        loss_test = loss_fn(out_test, test_labels)\n",
    "      else:\n",
    "        # (BONUS code) custom loss function from Bonus exercise 3.3\n",
    "        loss_test = loss_fn(out_test, test_labels, net.in_layer.weight,\n",
    "                            L2_penalty, L1_penalty)\n",
    "      test_loss.append(loss_test.item())  # .item() needed to transform the tensor output of loss_fn to a scalar\n",
    "\n",
    "    # Track progress\n",
    "    if (i + 1) % (n_iter // 5) == 0:\n",
    "      if test_data is None:\n",
    "        print(f'iteration {i + 1}/{n_iter} | loss: {loss.item():.3f}')\n",
    "      else:\n",
    "        print(f'iteration {i + 1}/{n_iter} | loss: {loss.item():.3f} | test_loss: {loss_test.item():.3f}')\n",
    "\n",
    "  if test_data is None:\n",
    "    return train_loss\n",
    "  else:\n",
    "    return train_loss, test_loss\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "# Initialize network with 10 hidden units\n",
    "net = DeepNetReLU(n_neurons, 10)\n",
    "\n",
    "# Initialize built-in PyTorch MSE loss function\n",
    "loss_fn = nn.MSELoss()\n",
    "\n",
    "# Run GD on data\n",
    "train_loss = train(net, loss_fn, resp_train, stimuli_train)\n",
    "\n",
    "# Plot the training loss over iterations of GD\n",
    "with plt.xkcd():\n",
    "  plt.plot(train_loss)\n",
    "  plt.xlim([0, None])\n",
    "  plt.ylim([0, None])\n",
    "  plt.xlabel('iterations of gradient descent')\n",
    "  plt.ylabel('mean squared error')\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 4: Investigating model and evaluating performance\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 4.1: Peering inside the decoding model\n",
    "\n",
    "We have built a model to perform decoding that takes as input neural activity and outputs the estimated angle of the stimulus. We can imagine that an animal that needs to determine angles would have a brain area that acts like the hidden layer in our model. It transforms the neural activity from visual cortex and outputs a decision. Decisions about orientations of edges could include figuring out how to jump onto a branch, how to avoid obstacles, or determining the type of an object, e.g. food or predator.\n",
    "\n",
    "What sort of connectivity would this brain area have with visual cortex? Determining this experimentally would be very difficult, perhaps we can look at the model we have and see if its structure constrains the type of connectivity we'd expect.\n",
    "\n",
    "Below we will visualize the weights from the neurons in visual cortex to the hidden units $W_{in}$, and the weights from the hidden units to the output orientation $W_{out}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:15.463270Z",
     "iopub.status.busy": "2021-05-31T20:48:15.447210Z",
     "iopub.status.idle": "2021-05-31T20:48:16.821824Z",
     "shell.execute_reply": "2021-05-31T20:48:16.821308Z"
    }
   },
   "outputs": [],
   "source": [
    "W_in = net.in_layer.weight.detach().numpy() # we can run .detach() and .numpy() to get a numpy array\n",
    "print('shape of W_in:')\n",
    "print(W_in.shape)\n",
    "\n",
    "W_out = net.out_layer.weight.detach().numpy() # we can run .detach() and .numpy() to get a numpy array\n",
    "print('shape of W_out:')\n",
    "print(W_out.shape)\n",
    "\n",
    "plt.figure(figsize=(10,4))\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(W_in, aspect='auto', cmap='bwr', vmin=-1e-2, vmax=1e-2)\n",
    "plt.xlabel('neurons')\n",
    "plt.ylabel('hidden units')\n",
    "plt.colorbar()\n",
    "plt.title('$W_{in}$')\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(W_out.T, cmap='bwr', vmin=-3, vmax=3)\n",
    "plt.xticks([])\n",
    "plt.xlabel('output')\n",
    "plt.ylabel('hidden units')\n",
    "plt.colorbar()\n",
    "plt.title('$W_{out}$')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coding Exercise 4.1: Visualizing weights\n",
    "\n",
    "It's difficult to see any structure in this weight matrix. How might we visualize it in a better way? \n",
    "\n",
    "Perhaps we can sort the neurons by their preferred orientation. We will use the `resp_all` matrix which is 360 stimuli (360$^\\circ$ of angles) by number of neurons. How do we find the preferred orientation? \n",
    "\n",
    "Let's visualize one column of this `resp_all` matrix first as we did at the beginning of the notebook. Can you see how we might want to first process this tuning curve before choosing the preferred orientation?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:16.841485Z",
     "iopub.status.busy": "2021-05-31T20:48:16.839582Z",
     "iopub.status.idle": "2021-05-31T20:48:17.041840Z",
     "shell.execute_reply": "2021-05-31T20:48:17.041385Z"
    }
   },
   "outputs": [],
   "source": [
    "idx = 235\n",
    "plt.plot(resp_all[:,idx])\n",
    "plt.ylabel('neural response')\n",
    "plt.xlabel('stimulus orientation ($^\\circ$)')\n",
    "plt.title(f'neuron {idx}')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looking at this tuning curve, there is a bit of noise across orientations, so let's smooth with a gaussian filter and then find the position of the maximum for each neuron. After getting the maximum position aka the \"preferred orientation\" for each neuron, we will re-sort the $W_{in}$ matrix. The maximum position in a matrix can be computed using the `.argmax(axis=_)` function in python -- make sure you specify the right axis though! Next, to get the indices of a matrix sorted you will need to use the `.argsort()` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:17.047966Z",
     "iopub.status.busy": "2021-05-31T20:48:17.047433Z",
     "iopub.status.idle": "2021-05-31T20:48:17.676387Z",
     "shell.execute_reply": "2021-05-31T20:48:17.675605Z"
    }
   },
   "outputs": [],
   "source": [
    "from scipy.ndimage import gaussian_filter1d\n",
    "\n",
    "# first let's smooth the tuning curves to make sure we get an accurate peak\n",
    "# that isn't just noise\n",
    "resp_smoothed = gaussian_filter1d(resp_all, 5, axis=0)\n",
    "\n",
    "###################################################\n",
    "## TO DO for students: find preferred orientation and resort W_in\n",
    "###################################################\n",
    "\n",
    "# find position of max response for each neuron\n",
    "# aka preferred orientation for each neuron\n",
    "preferred_orientation = ...\n",
    "\n",
    "# resort W_in matrix by preferred orientation\n",
    "isort = ...\n",
    "W_in_sorted = ...\n",
    "\n",
    "# plot resorted W_in matrix\n",
    "plt.figure(figsize=(10,4))\n",
    "plt.subplot(1,2,1)\n",
    "# plt.imshow(W_in_sorted, aspect='auto', cmap='bwr', vmin=-1e-2, vmax=1e-2)\n",
    "# plt.colorbar()\n",
    "plt.xlabel('sorted neurons')\n",
    "plt.ylabel('hidden units')\n",
    "plt.title('$W_{in}$')\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(W_out.T, cmap='bwr', vmin=-3, vmax=3)\n",
    "plt.xticks([])\n",
    "plt.xlabel('output')\n",
    "plt.ylabel('hidden units')\n",
    "plt.colorbar()\n",
    "plt.title('$W_{out}$')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:17.683620Z",
     "iopub.status.busy": "2021-05-31T20:48:17.683005Z",
     "iopub.status.idle": "2021-05-31T20:48:19.390428Z",
     "shell.execute_reply": "2021-05-31T20:48:19.389965Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "\n",
    "from scipy.ndimage import gaussian_filter1d\n",
    "\n",
    "# first let's smooth the tuning curves to make sure we get an accurate peak\n",
    "# that isn't just noise\n",
    "resp_smoothed = gaussian_filter1d(resp_all, 5, axis=0)\n",
    "\n",
    "###################################################\n",
    "## TO DO for students: find preferred orientation and resort W_in\n",
    "###################################################\n",
    "\n",
    "# find position of max response for each neuron\n",
    "# aka preferred orientation for each neuron\n",
    "preferred_orientation = resp_smoothed.argmax(axis=0)\n",
    "\n",
    "# resort W_in matrix by preferred orientation\n",
    "isort = preferred_orientation.argsort()\n",
    "W_in_sorted = W_in[:,isort]\n",
    "\n",
    "# plot resorted W_in matrix\n",
    "with plt.xkcd():\n",
    "  plt.figure(figsize=(10,4))\n",
    "  plt.subplot(1,2,1)\n",
    "  plt.imshow(W_in_sorted, aspect='auto', cmap='bwr', vmin=-1e-2, vmax=1e-2)\n",
    "  plt.xlabel('sorted neurons')\n",
    "  plt.ylabel('hidden units')\n",
    "  plt.colorbar()\n",
    "  plt.title('$W_{in}$')\n",
    "\n",
    "  plt.subplot(1,2,2)\n",
    "  plt.imshow(W_out.T, cmap='bwr', vmin=-3, vmax=3)\n",
    "  plt.xticks([])\n",
    "  plt.xlabel('output')\n",
    "  plt.ylabel('hidden units')\n",
    "  plt.colorbar()\n",
    "  plt.title('$W_{out}$')\n",
    "\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Think! 4: Interpreting weights\n",
    "\n",
    "We have just visualized how the model transforms neural activity to hidden layer activity. How should we interpret these matrices? Here are some guiding questions to explore:\n",
    "* Why are some of the $W_{in}$ weights close to zero for some of the hidden units? Do these correspond to close to zero weights in $W_{out}$?\n",
    "* Note how each hidden unit seems to have strongest weights to two groups of neurons in $W_{in}$, corresponding to two different sets of preferred orientations. Why do you think that is? What does this tell us about the structure of the tuning curves of the neurons?\n",
    "* How can we better understand the hidden units? What simple visualization would you try next?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:19.395374Z",
     "iopub.status.busy": "2021-05-31T20:48:19.394781Z",
     "iopub.status.idle": "2021-05-31T20:48:19.397344Z",
     "shell.execute_reply": "2021-05-31T20:48:19.397742Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove explanation\n",
    "\n",
    "\"\"\"\n",
    "1. It seems like the model did not learn any weights for those hidden units,\n",
    "perhaps the problem is simple enough that the weights were not needed, or the\n",
    "random initialization of the W_out weights for those hidden units close to zero\n",
    "kept them from changing. You could test the second hypothesis by looking at the\n",
    "weights at initialization\n",
    "2. Neurons often have tuning preference to two gratings at 180 degrees apart\n",
    "since these two gratings are the same other than the phase. This is likely why\n",
    "there are two peaks.\n",
    "3. One could look at the activity of the hidden units across various stimuli,\n",
    "making tuning curves for the hidden units.\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 4.2: Generalization performance with test data\n",
    "\n",
    "Note that gradient descent is essentially an algorithm for fitting the network's parameters to a given set of training data. Selecting this training data is thus crucial for ensuring that the optimized parameters **generalize** to unseen data they weren't trained on. In our case, for example, we want to make sure that our trained network is good at decoding stimulus orientations from neural responses to any orientation, not just those in our data set.\n",
    "\n",
    "To ensure this, we have split up the full data set into a **training set** and a **testing set**. In Coding Exercise 3.2, we trained a deep network by optimizing the parameters on a training set. We will now evaluate how good the optimized parameters are by using the trained network to decode stimulus orientations from neural responses in the testing set. Good decoding performance on this testing set should then be indicative of good decoding performance on the neurons' responses to any other stimulus orientation. This procedure is commonly used in machine learning (not just in deep learning)and is typically referred to as **cross-validation**.\n",
    "\n",
    "We will compute the MSE on the test data and plot the decoded stimulus orientations as a function of the true stimulus.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:19.402746Z",
     "iopub.status.busy": "2021-05-31T20:48:19.402222Z",
     "iopub.status.idle": "2021-05-31T20:48:19.598855Z",
     "shell.execute_reply": "2021-05-31T20:48:19.598410Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "#@markdown Execute this cell to evaluate and plot test error\n",
    "\n",
    "out = net(resp_test)  # decode stimulus orientation for neural responses in testing set\n",
    "ori = stimuli_test  # true stimulus orientations\n",
    "test_loss = loss_fn(out, ori)  # MSE on testing set (Hint: use loss_fn initialized in previous exercise)\n",
    "\n",
    "plt.plot(ori, out.detach(), '.')  # N.B. need to use .detach() to pass network output into plt.plot()\n",
    "identityLine()  # draw the identity line y=x; deviations from this indicate bad decoding!\n",
    "plt.title('MSE on testing set: %.2f' % test_loss.item())  # N.B. need to use .item() to turn test_loss into a scalar\n",
    "plt.xlabel('true stimulus orientation ($^o$)')\n",
    "plt.ylabel('decoded stimulus orientation ($^o$)')\n",
    "axticks = np.linspace(0, 360, 5)\n",
    "plt.xticks(axticks)\n",
    "plt.yticks(axticks)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**PyTorch Note**:\n",
    "\n",
    "An important thing to note in the code snippet for plotting the decoded orientations is the `.detach()` method. The PyTorch `nn.Module` class is special in that, behind the scenes, each of the variables inside it are linked to each other in a computational graph, for the purposes of automatic differentiation (the algorithm used in `.backward()` to compute gradients). As a result, if you want to do anything that is not a `torch` operation to the parameters or outputs of an `nn.Module` class, you'll need to first \"detach\" it from its computational graph. This is what the `.detach()` method does. In this hidden code above, we need to call it on the outputs of the network so that we can plot them with the `plt.plot()` function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If interested, please see Bonus Section 3 to think more about model criticism (3.1), improve the loss function accordingly (3.2), and add regularization (3.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Summary\n",
    "\n",
    "We have now covered a number of common and powerful techniques for applying deep learning to decoding from neural data, some of which are common to almost any machine learning problem:\n",
    "* Building and training deep networks using the **PyTorch** `nn.Module` class and built-in **optimizers**\n",
    "* Choosing and evaluating **loss functions**\n",
    "* Testing a trained model on unseen data via **cross-validation**, by splitting the data into a **training set and testing set**\n",
    "\n",
    "An important aspect of this tutorial was the `train()` function we wrote in exercise 3. Note that it can be used to train *any* network to minimize *any* loss function (cf. advanced exercise 6) on *any* training data. This is the power of using PyTorch to train neural networks and, for that matter, **any other model**! There is nothing in the `nn.Module` class that forces us to use `nn.Linear` layers that implement neural network operations. You can actually put anything you want inside the `.__init__()` and `.forward()` methods of this class. As long as its parameters and computations involve only `torch.Tensor`'s, and the model is differentiable, you'll then be able to optimize the parameters of this model in exactly the same way we optimized the deep networks here.\n",
    "\n",
    "What kinds of conclusions can we draw from these sorts of analyses? If we can decode the stimulus well from visual cortex activity, that means that there is information about this stimulus available in visual cortex. Whether or not the animal uses that information to make decisions is not determined from an analysis like this. In fact mice perform poorly in orientation discrimination tasks compared to monkeys and humans, even though they have information about these stimuli in their visual cortex. Why do you think they perform poorly in orientation discrimination tasks?\n",
    "\n",
    "See [this paper](https://www.biorxiv.org/content/10.1101/679324v2) for some potential hypotheses, but this is totally an open question!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Bonus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus Section 1: Neural network *depth*, *width* and *expressivity*\n",
    "\n",
    "Two important architectural choices that always have to be made when constructing deep feed-forward networks like those used here are\n",
    "* the number of hidden layers, or the network's *depth*\n",
    "* the number of units in each layer, or the layer *widths*\n",
    "\n",
    "Here, we restricted ourselves to networks with a single hidden layer with a width of $M$ units, but it is easy to see how this code could be adapted to arbitrary depths. Adding another hidden layer simply requires adding another `nn.Linear` module to the `__init__()` method and incorporating it into the `.forward()` method.\n",
    "\n",
    "The depth and width of a network determine the set of input/output transormations that it can perform, often referred to as its *expressivity*. The deeper and wider the network, the more *expressive* it is; that is, the larger the class of input/output transformations it can compute. In fact, it turns out that an infinitely wide *or* infinitely deep networks can in principle [compute (almost) *any* input/output transformation](https://en.wikipedia.org/wiki/Universal_approximation_theorem).\n",
    "\n",
    "A classic mathematical demonstration of the power of depth is given by the so-called [XOR problem](https://medium.com/@jayeshbahire/the-xor-problem-in-neural-networks-50006411840b#:~:text=The%20XOr%2C%20or%20%E2%80%9Cexclusive%20or,value%20if%20they%20are%20equal.). This toy problem demonstrates how even a single hidden layer can drastically expand the set of input/output transformations a network can perform, relative to a shallow network with no hidden layers. The key intuition is that the hidden layer allows you to represent the input in a new format, which can then allow you to do almost anything you want with it. The *wider* this hidden layer, the more flexibility you have in this representation. In particular, if you have more hidden units than input units, then the hidden layer representation of the input is higher-dimensional than the raw data representation. This higher dimensionality effectively gives you more \"room\" to perform arbitrary computations in. It turns out that even with just this one hidden layer, if you make it wide enough you can actually approximate any input/output transformation you want. See [here](http://neuralnetworksanddeeplearning.com/chap4.html) for a neat visual demonstration of this.\n",
    "\n",
    "In practice, however, it turns out that increasing depth seems to grant more expressivity with fewer units than increasing width does (for reasons that are not well understood). It is for this reason that truly *deep* networks are almost always used in machine learning, which is why this set of techniques is often referred to as *deep* learning.\n",
    "\n",
    "That said, there is a cost to making networks deeper and wider. The bigger your network, the more parameters (i.e. weights and biases) it has, which need to be optimized! The extra expressivity afforded by higher width and/or depth thus carries with it (at least) two problems:\n",
    "* optimizing more parameters usually requires more data\n",
    "* a more highly parameterized network is more prone to overfit to the training data, so requires more sophisticated optimization algorithms to ensure generalization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus Section 2: Gradient descent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus Section 2.1: Gradient descent equations\n",
    "\n",
    "Here we provide the equations for the three steps of the gradient descent algorithm, as applied to our decoding problem:\n",
    "\n",
    "1. **Evaluate the loss** on the training data. For a mean squared error loss, this is given by\n",
    "\\begin{equation}\n",
    "    L = \\frac{1}{P}\\sum_{n=1}^P (y^{(n)} - \\tilde{y}^{(n)})^2\n",
    "\\end{equation}\n",
    "where $y^{(n)}$ denotes the stimulus orientation decoded from the population response $\\mathbf{r}^{(n)}$ to the $n$th stimulus in the training data, and $\\tilde{y}^{(n)}$ is the true orientation of that stimulus. $P$ denotes the total number of data samples in the training set. In the syntax of our `train()` function above, $\\mathbf{r}^{(n)}$ is given by `train_data[n, :]` and $\\tilde{y}^{(n)}$ by `train_labels[n]`.\n",
    "\n",
    "2. **Compute the gradient of the loss** with respect to each of the network weights. In our case, this entails computing the quantities\n",
    "\\begin{equation}\n",
    "    \\frac{\\partial L}{\\partial \\mathbf{W}^{in}}, \\frac{\\partial L}{\\partial \\mathbf{b}^{in}}, \\frac{\\partial L}{\\partial \\mathbf{W}^{out}}, \\frac{\\partial L}{\\partial \\mathbf{b}^{out}}\n",
    "\\end{equation}\n",
    "Usually, we would require lots of math in order to derive each of these gradients, and lots of code to compute them. But this is where PyTorch comes to the rescue! Using a cool technique called [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation), PyTorch automatically calculates these gradients when the `.backward()` function is called.\n",
    "\n",
    "  More specifically, when this function is called on a particular variable (e.g. `loss`, as above), PyTorch will compute the gradients with respect to each network parameter. These are computed and stored behind the scenes, and can be accessed through the `.grad` attribute of each of the network's parameters. As we saw above, however, we actually never need to look at or call these gradients when implementing gradient descent, as this can be taken care of by PyTorch's built-in optimizers, like `optim.SGD`.\n",
    "\n",
    "3. **Update the network weights** by descending the gradient:\n",
    "\\begin{align}\n",
    "    \\mathbf{W}^{in} &\\leftarrow \\mathbf{W}^{in} - \\alpha \\frac{\\partial L}{\\partial \\mathbf{W}^{in}} \\\\\n",
    "    \\mathbf{b}^{in} &\\leftarrow \\mathbf{b}^{in} - \\alpha \\frac{\\partial L}{\\partial \\mathbf{b}^{in}} \\\\\n",
    "    \\mathbf{W}^{out} &\\leftarrow \\mathbf{W}^{out} - \\alpha \\frac{\\partial L}{\\partial \\mathbf{W}^{out}} \\\\\n",
    "    \\mathbf{b}^{out} &\\leftarrow \\mathbf{b}^{out} - \\alpha \\frac{\\partial L}{\\partial \\mathbf{b}^{out}}\n",
    "\\end{align}\n",
    "where $\\alpha$ is called the **learning rate**. This **hyperparameter** of the SGD algorithm controls how far we descend the gradient on each iteration. It should be as large as possible so that fewer iterations are needed, but not too large so as to avoid parameter updates from skipping over minima in the loss landscape.\n",
    "\n",
    "While the equations written down here are specific to the network and loss function considered in this tutorial, the code provided above for implementing these three steps is completely general: no matter what loss function or network you are using, exactly the same commands can be used to implement these three steps.\n",
    "\n",
    "The way that the gradients are calculated is called **backpropagation**. We have a loss function:\n",
    "\\begin{align}\n",
    "    L &= (y - \\tilde{y})^2 \\\\\n",
    "     &= (\\mathbf{W}^{out} \\mathbf{h} - \\tilde{y})^2\n",
    "\\end{align}\n",
    "where $\\mathbf{h} = \\phi(\\mathbf{W}^{in} \\mathbf{r} + \\mathbf{b}^{in})$\n",
    "You may see that $\\frac{\\partial L}{\\partial \\mathbf{W}^{out}}$ is simple to calculate as it is on the outside of the equation (it is also a vector in this case, not a matrix, so the derivative is standard):\n",
    "\\begin{equation}\n",
    "\\frac{\\partial L}{\\partial \\mathbf{W}^{out}} = 2 (\\mathbf{h} - \\tilde{y})\n",
    "\\end{equation}\n",
    "Now let's compute the derivative with respect to $\\mathbf{W}^{in}$ using the chain rule. Note it is only positive if the output is positive due to the RELU $\\phi$:\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial \\mathbf{W}^{in}} &= \\begin{cases}\n",
    "    \\frac{\\partial L}{\\partial \\mathbf{W}^{out}} \\frac{\\partial \\mathbf{h}}{\\partial \\mathbf{W}^{in}}  & \\text{if }  \\mathbf{h} > 0 \\\\\n",
    "    0 & \\text{else}\n",
    "  \\end{cases} \\\\\n",
    "  &= \\begin{cases}\n",
    "    2 (\\mathbf{h} - \\tilde{y}) \\mathbf{r}^\\top  & \\text{if }  \\mathbf{h} > 0 \\\\\n",
    "    0 & \\text{else}\n",
    "  \\end{cases}\n",
    "\\end{align}\n",
    "It is most efficient to compute the derivative once for the last layer, then once for the next layer and multiply by the previous layer's derivative and so on using the chain rule. Each of these operations is relatively fast, making training of deep networks feasible. \n",
    "\n",
    "The command `loss.backward()` computes these gradients for the defined `loss` with respect to each network parameter. The computation is done using [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation), which implements backpropagation. Note that this works no matter how big/small the network is, allowing us to perform gradient descent for any deep network model built using PyTorch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus Section 2.2: *Stochastic* gradient descent (SGD) vs. gradient descent (GD)\n",
    "\n",
    "In this tutorial, we used the gradient descent algorithm, which differs in a subtle yet very important way from the more commonly used **stochastic gradient descent (SGD)** algorithm. The key difference is in the very first step of each iteration, where in the GD algorithm we evaluate the loss *at every data sample in the training set*. In SGD, on the other hand, we evaluate the loss only at a random subset of data samlpes from the full training set, called a **mini-batch**. At each iteration, we randomly sample a mini-batch to perform steps 1-3 on. All the above equations still hold, but now the $P$ data samples $\\mathbf{r}^{(n)}, \\tilde{y}^{(n)}$ denote a mini-batch of $P$ random samples from the training set, rather than the whole training set.\n",
    "\n",
    "There are several reasons why one might want to use SGD instead of GD. The first is that the training set might be too big, so that we actually can't actually evaluate the loss on every single data sample in it. In this case, GD is simply infeasible, so we have no choice but to turn to SGD, which bypasses the restrictive memory demands of GD by sub-sampling the training set into smaller mini-batches.\n",
    "\n",
    "But, even when GD is feasible, SGD turns out to be generally better. The stochasticity induced by the extra random sampling step in SGD effectively adds some noise in the search for local minima of the loss function. This can be really useful for avoiding potential local minima, and enforce that whatever minimum is converged to is a good one. This is particularly important when networks are wider and/or deeper, in which case the large number of parameters can lead to overfitting.\n",
    "\n",
    "Here, we used only GD because (1) it is simpler, and (2) it suffices for the problem being considered here. Because we have so many neurons in our data set, decoding is not too challenging and doesn't require a particularly deep or wide network. The small number of parameters in our deep networks therefore can be optimized without a problem using GD."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus Section 3: Evaluating & improving models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "### Bonus Section 3.1: Model criticism\n",
    "\n",
    "Let's now take a step back and think about how our model is succeeding/failing and how to improve it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:19.608069Z",
     "iopub.status.busy": "2021-05-31T20:48:19.607521Z",
     "iopub.status.idle": "2021-05-31T20:48:19.790943Z",
     "shell.execute_reply": "2021-05-31T20:48:19.791491Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "#@markdown Execute this cell to plot decoding error\n",
    "\n",
    "out = net(resp_test)  # decode stimulus orientation for neural responses in testing set\n",
    "ori = stimuli_test  # true stimulus orientations\n",
    "error = out - ori  # decoding error\n",
    "\n",
    "\n",
    "plt.plot(ori, error.detach(), '.')   # plot decoding error as a function of true orientation (make sure all arguments to plt.plot() have been detached from PyTorch network!)\n",
    "\n",
    "# Plotting\n",
    "plt.xlabel('true stimulus orientation ($^o$)')\n",
    "plt.ylabel('decoding error ($^o$)')\n",
    "plt.xticks(np.linspace(0, 360, 5))\n",
    "plt.yticks(np.linspace(-360, 360, 9))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Bonus Think! 3.1: Delving into error problems\n",
    "\n",
    "In the cell below, we will plot the *decoding error* for each neural response in the testing set. The decoding error is defined as the decoded stimulus orientation minus true stimulus orientation\n",
    "\\begin{equation}\n",
    "  \\text{decoding error} = y^{(n)} - \\tilde{y}^{(n)}\n",
    "\\end{equation}\n",
    "\n",
    "In particular, we plot decoding error as a function of the true stimulus orientation.\n",
    "\n",
    "\n",
    "  * Are some stimulus orientations harder to decode than others?\n",
    "  * If so, in what sense? Are the decoded orientations for these stimuli more variable and/or are they biased?\n",
    "  * Can you explain this variability/bias? What makes these stimulus orientations different from the others?\n",
    "  * (Will be addressed in next exercise) Can you think of a way to modify the deep network in order to avoid this?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:19.797016Z",
     "iopub.status.busy": "2021-05-31T20:48:19.796040Z",
     "iopub.status.idle": "2021-05-31T20:48:19.799049Z",
     "shell.execute_reply": "2021-05-31T20:48:19.799535Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove explanation\n",
    "\n",
    "\"\"\"\n",
    "It appears that the errors are larger at 0 and 360 degrees. The errors are biased\n",
    "in the positive direction at 0 degrees and in the negative direction at 360 degrees.\n",
    "This is because the 0 degree stimulus and the 360 degree stimulus are in fact the\n",
    "same because orientation is a circular variable. The network therefore has trouble\n",
    "determining whether the stimulus is 0 or 360 degrees.\n",
    "\n",
    "We can modify the deep network to avoid this problem in a few different ways.\n",
    "One approach would be to predict a sine and a cosine of the angle and then taking\n",
    " the predicted angle as the angle of the complex number $sin(\\theta) + i cos(\\theta)$.\n",
    "\n",
    "An alternative approach is to bin the stimulus responses and predict the bin of the stimulus.\n",
    "This turns the problem into a classification problem rather than a regression problem,\n",
    "and in this case you will need to use a new loss function (see below).\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus Section 3.2: Improving the loss function \n",
    "As illustrated in the previous exercise, the squared error is not a good loss function for circular quantities like angles, since two angles that are very close (e.g. $1^o$ and $359^o$) might actually have a very large squared error.\n",
    "\n",
    "Here, we'll avoid this problem by changing our loss function to treat our decoding problem as a **classification problem**. Rather than estimating the *exact* angle of the stimulus, we'll now aim to construct a decoder that classifies the stimulus into one of $C$ classes, corresponding to different bins of angles of width $b = \\frac{360}{C}$. The true class $\\tilde{y}^{(n)}$ of stimulus $i$ is now given by\n",
    "\\begin{equation}\n",
    "  \\tilde{y}^{(n)} =\n",
    "  \\begin{cases}\n",
    "    1 &\\text{if angle of stimulus $n$ is in the range } [0, b] \\\\\n",
    "    2 &\\text{if angle of stimulus $n$ is in the range } [b, 2b] \\\\\n",
    "    3 &\\text{if angle of stimulus $n$ is in the range } [2b, 3b] \\\\\n",
    "    \\vdots \\\\\n",
    "    C &\\text{if angle of stimulus $n$ is in the range } [(C-1)b, 360]\n",
    "  \\end{cases}\n",
    "\\end{equation}\n",
    "\n",
    "We have a helper function `stimulus_class` that will extract `n_classes` stimulus classes for us from the stimulus orientations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To decode the stimulus class from neural responses, we'll use a deep network that outputs a $C$-dimensional vector of probabilities $\\mathbf{p} = \\begin{bmatrix} p_1, p_2, \\ldots, p_C \\end{bmatrix}^T$, corresponding to the estimated probabilities of the stimulus belonging to each class $1, 2, \\ldots, C$. \n",
    "\n",
    "To ensure the network's outputs are indeed probabilities (i.e. they are positive numbers between 0 and 1, and sum to 1), we'll use a [softmax function](https://en.wikipedia.org/wiki/Softmax_function) to transform the real-valued outputs from the hidden layer into probabilities. Letting $\\sigma(\\cdot)$ denote this softmax function, the equations describing our network are\n",
    "\\begin{align}\n",
    "    \\mathbf{h}^{(n)} &= \\phi(\\mathbf{W}^{in} \\mathbf{r}^{(n)} + \\mathbf{b}^{in}), && [\\mathbf{W}^{in}: M \\times N], \\\\\n",
    "    \\mathbf{p}^{(n)} &= \\sigma(\\mathbf{W}^{out} \\mathbf{h}^{(n)} + \\mathbf{b}^{out}),  && [\\mathbf{W}^{out}: C \\times M],\n",
    "\\end{align}\n",
    "The decoded stimulus class is then given by that assigned the highest probability by the network:\n",
    "\\begin{equation}\n",
    "  y^{(n)} = \\underset{i}{\\arg\\max} \\,\\, p_i\n",
    "\\end{equation}\n",
    "The softmax function can be implemented in PyTorch simply using `torch.softmax()`.\n",
    "\n",
    "Often *log* probabilities are easier to work with than actual probabilities, because probabilities tend to be very small numbers that computers have trouble representing. We'll therefore actually use the logarithm of the softmax as the output of our network,\n",
    "\\begin{equation}\n",
    "    \\mathbf{l}^{(n)} = \\log \\left( \\mathbf{p}^{(n)} \\right)\n",
    "\\end{equation}\n",
    "which can implemented in PyTorch together with the softmax via an `nn.LogSoftmax` layer. The nice thing about the logarithmic function is that it's *monotonic*, so if one probability is larger/smaller than another, then its logarithm is also larger/smaller than the other's. We therefore have that\n",
    "\\begin{equation}\n",
    "  y^{(n)} = \\underset{i}{\\arg\\max} \\,\\, p_i^{(n)} = \\underset{i}{\\arg\\max} \\, \\log p_i^{(n)} = \\underset{i}{\\arg\\max} \\,\\, l_i^{(n)}\n",
    "\\end{equation}\n",
    "\n",
    "See the next cell for code for constructing a deep network with one hidden layer that of ReLU's that outputs a vector of log probabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:19.806258Z",
     "iopub.status.busy": "2021-05-31T20:48:19.805046Z",
     "iopub.status.idle": "2021-05-31T20:48:19.806892Z",
     "shell.execute_reply": "2021-05-31T20:48:19.807374Z"
    }
   },
   "outputs": [],
   "source": [
    "# Deep network for classification\n",
    "class DeepNetSoftmax(nn.Module):\n",
    "  \"\"\"Deep Network with one hidden layer, for classification\n",
    "\n",
    "  Args:\n",
    "    n_inputs (int): number of input units\n",
    "    n_hidden (int): number of units in hidden layer\n",
    "    n_classes (int): number of outputs, i.e. number of classes to output\n",
    "      probabilities for\n",
    "\n",
    "  Attributes:\n",
    "    in_layer (nn.Linear): weights and biases of input layer\n",
    "    out_layer (nn.Linear): weights and biases of output layer\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  def __init__(self, n_inputs, n_hidden, n_classes):\n",
    "    super().__init__()  # needed to invoke the properties of the parent class nn.Module\n",
    "    self.in_layer = nn.Linear(n_inputs, n_hidden)  # neural activity --> hidden units\n",
    "    self.out_layer = nn.Linear(n_hidden, n_classes)  # hidden units --> outputs\n",
    "    self.logprob = nn.LogSoftmax(dim=1)  # probabilities across columns should sum to 1 (each output row corresponds to a different input)\n",
    "\n",
    "  def forward(self, r):\n",
    "    \"\"\"Predict stimulus orientation bin from neural responses\n",
    "\n",
    "    Args:\n",
    "      r (torch.Tensor): n_stimuli x n_inputs tensor with neural responses to n_stimuli\n",
    "\n",
    "    Returns:\n",
    "      torch.Tensor: n_stimuli x n_classes tensor with predicted class probabilities\n",
    "\n",
    "    \"\"\"\n",
    "    h = torch.relu(self.in_layer(r))\n",
    "    logp = self.logprob(self.out_layer(h))\n",
    "    return logp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What should our loss function now be? Ideally, we want the probabilities outputted by our network to be such that the probability of the true stimulus class is high. One way to formalize this is to say that we want to maximize the *log* probability of the true stimulus class $\\tilde{y}^{(n)}$ under the class probabilities predicted by the network,\n",
    "\\begin{equation}\n",
    "  \\log \\left( \\text{predicted probability of stimulus } n \\text{ being of class } \\tilde{y}^{(n)} \\right) = \\log p^{(n)}_{\\tilde{y}^{(n)}} = l^{(n)}_{\\tilde{y}^{(n)}}\n",
    "\\end{equation}\n",
    "To turn this into a loss function to be *minimized*, we can then simply multiply it by -1: maximizing the log probability is the same as minimizing the *negative* log probability. Summing over a batch of $P$ inputs, our loss function is then given by\n",
    "\\begin{equation}\n",
    "  L = -\\sum_{n=1}^P \\log p^{(n)}_{\\tilde{y}^{(n)}} = -\\sum_{n=1}^P l^{(n)}_{\\tilde{y}^{(n)}}\n",
    "\\end{equation}\n",
    "In the deep learning community, this loss function is typically referred to as the **cross-entropy**, or **negative log likelihood**. The corresponding built-in loss function in PyTorch is `nn.NLLLoss()` (documentation [here](https://pytorch.org/docs/master/generated/torch.nn.CrossEntropyLoss.html)).\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Bonus Coding Exercise 3.2: A new loss function\n",
    "In the next cell, we've provided most of the code to train and test a network to decode stimulus orientations via classification, by minimizing the negative log likelihood. Fill in the missing pieces.\n",
    "\n",
    "Once you've done this, have a look at the plotted results. Does changing the loss function from mean squared error to a classification loss solve our problems? Note that errors may still occur -- but are these errors as bad as the ones that our network above was making?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:19.815623Z",
     "iopub.status.busy": "2021-05-31T20:48:19.814965Z",
     "iopub.status.idle": "2021-05-31T20:48:19.817701Z",
     "shell.execute_reply": "2021-05-31T20:48:19.817242Z"
    }
   },
   "outputs": [],
   "source": [
    "def decode_orientation(net, n_classes, loss_fn,\n",
    "                       train_data, train_labels, test_data, test_labels,\n",
    "                       n_iter=1000, L2_penalty=0, L1_penalty=0):\n",
    "  \"\"\" Initialize, train, and test deep network to decode binned orientation from neural responses\n",
    "\n",
    "  Args:\n",
    "    net (nn.Module): deep network to run\n",
    "    n_classes (scalar): number of classes in which to bin orientation\n",
    "    loss_fn (function): loss function to run\n",
    "    train_data (torch.Tensor): n_train x n_neurons tensor with neural\n",
    "      responses to train on\n",
    "    train_labels (torch.Tensor): n_train x 1 tensor with orientations of the\n",
    "      stimuli corresponding to each row of train_data, in radians\n",
    "    test_data (torch.Tensor): n_test x n_neurons tensor with neural\n",
    "      responses to train on\n",
    "    test_labels (torch.Tensor): n_test x 1 tensor with orientations of the\n",
    "      stimuli corresponding to each row of train_data, in radians\n",
    "    n_iter (int, optional): number of iterations to run optimization\n",
    "    L2_penalty (float, optional): l2 penalty regularizer coefficient\n",
    "    L1_penalty (float, optional): l1 penalty regularizer coefficient\n",
    "\n",
    "  Returns:\n",
    "    (list, torch.Tensor): training loss over iterations, n_test x 1 tensor with predicted orientations of the\n",
    "      stimuli from decoding neural network\n",
    "  \"\"\"\n",
    "\n",
    "  # Bin stimulus orientations in training set\n",
    "  train_binned_labels = stimulus_class(train_labels, n_classes)\n",
    "  test_binned_labels = stimulus_class(test_labels, n_classes)\n",
    "\n",
    "\n",
    "  # Run GD on training set data, using learning rate of 0.1\n",
    "  # (add optional arguments test_data and test_binned_labels!)\n",
    "  train_loss, test_loss = train(net, loss_fn, train_data, train_binned_labels,\n",
    "                                learning_rate=0.1, test_data=test_data,\n",
    "                                test_labels=test_binned_labels, n_iter=n_iter,\n",
    "                                L2_penalty=L2_penalty, L1_penalty=L1_penalty)\n",
    "\n",
    "  # Decode neural responses in testing set data\n",
    "  out = net(test_data)\n",
    "  out_labels = np.argmax(out.detach(), axis=1)  # predicted classes\n",
    "\n",
    "  frac_correct = (out_labels==test_binned_labels).sum() / len(test_binned_labels)\n",
    "  print(f'>>> fraction correct = {frac_correct:.3f}')\n",
    "\n",
    "  return train_loss, test_loss, out_labels\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "n_classes = 20\n",
    "\n",
    "##############################################################################\n",
    "## TODO for students: fill out missing pieces below to initialize network and\n",
    "# define loss function\n",
    "##############################################################################\n",
    "\n",
    "# Initialize network\n",
    "net = ... # use M=20 hidden units\n",
    "\n",
    "# Initialize built-in PyTorch negative log likelihood loss function\n",
    "loss_fn = ...\n",
    "\n",
    "# Uncomment below to train network and run it on test images\n",
    "# this function uses the train function you wrote before\n",
    "#train_loss, test_loss, predicted_test_labels = decode_orientation(net, n_classes, loss_fn,\n",
    "#                                                                  resp_train, stimuli_train, resp_test, stimuli_test)\n",
    "\n",
    "# Plot results\n",
    "#plot_decoded_results(train_loss, test_loss, stimuli_test, predicted_test_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:19.825918Z",
     "iopub.status.busy": "2021-05-31T20:48:19.825333Z",
     "iopub.status.idle": "2021-05-31T20:48:27.862160Z",
     "shell.execute_reply": "2021-05-31T20:48:27.861700Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "\n",
    "def decode_orientation(net, n_classes, loss_fn,\n",
    "                       train_data, train_labels, test_data, test_labels,\n",
    "                       n_iter=1000, L2_penalty=0, L1_penalty=0):\n",
    "  \"\"\" Initialize, train, and test deep network to decode binned orientation from neural responses\n",
    "\n",
    "  Args:\n",
    "    net (nn.Module): deep network to run\n",
    "    n_classes (scalar): number of classes in which to bin orientation\n",
    "    loss_fn (function): loss function to run\n",
    "    train_data (torch.Tensor): n_train x n_neurons tensor with neural\n",
    "      responses to train on\n",
    "    train_labels (torch.Tensor): n_train x 1 tensor with orientations of the\n",
    "      stimuli corresponding to each row of train_data, in radians\n",
    "    test_data (torch.Tensor): n_test x n_neurons tensor with neural\n",
    "      responses to train on\n",
    "    test_labels (torch.Tensor): n_test x 1 tensor with orientations of the\n",
    "      stimuli corresponding to each row of train_data, in radians\n",
    "    n_iter (int, optional): number of iterations to run optimization\n",
    "    L2_penalty (float, optional): l2 penalty regularizer coefficient\n",
    "    L1_penalty (float, optional): l1 penalty regularizer coefficient\n",
    "\n",
    "  Returns:\n",
    "    (list, torch.Tensor): training loss over iterations, n_test x 1 tensor with predicted orientations of the\n",
    "      stimuli from decoding neural network\n",
    "  \"\"\"\n",
    "\n",
    "  # Bin stimulus orientations in training set\n",
    "  train_binned_labels = stimulus_class(train_labels, n_classes)\n",
    "  test_binned_labels = stimulus_class(test_labels, n_classes)\n",
    "\n",
    "\n",
    "  # Run GD on training set data, using learning rate of 0.1\n",
    "  # (add optional arguments test_data and test_binned_labels!)\n",
    "  train_loss, test_loss = train(net, loss_fn, train_data, train_binned_labels,\n",
    "                                learning_rate=0.1, test_data=test_data,\n",
    "                                test_labels=test_binned_labels, n_iter=n_iter,\n",
    "                                L2_penalty=L2_penalty, L1_penalty=L1_penalty)\n",
    "\n",
    "  # Decode neural responses in testing set data\n",
    "  out = net(test_data)\n",
    "  out_labels = np.argmax(out.detach(), axis=1)  # predicted classes\n",
    "\n",
    "  frac_correct = (out_labels==test_binned_labels).sum() / len(test_binned_labels)\n",
    "  print(f'>>> fraction correct = {frac_correct:.3f}')\n",
    "\n",
    "  return train_loss, test_loss, out_labels\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "n_classes = 20\n",
    "\n",
    "# Initialize network\n",
    "net = DeepNetSoftmax(n_neurons, 20, n_classes)  # use M=20 hidden units\n",
    "\n",
    "# Initialize built-in PyTorch negative log likelihood loss function\n",
    "loss_fn = nn.NLLLoss()\n",
    "\n",
    "# Uncomment below to train network and run it on test images\n",
    "# this function uses the train function you wrote before\n",
    "train_loss, test_loss, predicted_test_labels = decode_orientation(net, n_classes, loss_fn,\n",
    "                                                                  resp_train, stimuli_train, resp_test, stimuli_test)\n",
    "\n",
    "# Plot results\n",
    "with plt.xkcd():\n",
    "  plot_decoded_results(train_loss, test_loss, stimuli_test, predicted_test_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How do the weights $W_{in}$ from the neurons to the hidden layer look now?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:27.868441Z",
     "iopub.status.busy": "2021-05-31T20:48:27.867978Z",
     "iopub.status.idle": "2021-05-31T20:48:28.980030Z",
     "shell.execute_reply": "2021-05-31T20:48:28.980422Z"
    }
   },
   "outputs": [],
   "source": [
    "W_in = net.in_layer.weight.detach().numpy() # we can run detach and numpy to get a numpy array\n",
    "print('shape of W_in:')\n",
    "print(W_in.shape)\n",
    "\n",
    "plt.figure(figsize=(6,4))\n",
    "plt.imshow(W_in[:,isort], aspect='auto', cmap='bwr', vmin=-5e-3, vmax=5e-3)\n",
    "plt.xlabel('neurons')\n",
    "plt.ylabel('hidden units')\n",
    "plt.colorbar()\n",
    "plt.title('$W_{in}$')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus Section 3.3: Regularization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:28.985652Z",
     "iopub.status.busy": "2021-05-31T20:48:28.985077Z",
     "iopub.status.idle": "2021-05-31T20:48:28.998514Z",
     "shell.execute_reply": "2021-05-31T20:48:28.998981Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 4: Regularization\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"Qnn5OPHKo5w\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "As discussed in the lecture, it is often important to incorporate regularization terms into the loss function to avoid overfitting. In particular, in this case, we will use these terms to enforce sparsity in the linear layer from neurons to hidden units. \n",
    "\n",
    "Here we'll consider the classic L2 regularization penalty $\\mathcal{R}_{L2}$, which is the sum of squares of each weight in the network $\\sum_{ij} {\\mathbf{W}^{out}_{ij}}^2$ times a constant that we call `L2_penalty`.\n",
    "\n",
    "We will also add an L1 regularization penalty $\\mathcal{R}_{L1}$ to enforce sparsity of the weights, which is the sum of the absolute values of the weights $\\sum_{ij} |{\\mathbf{W}^{out}_{ij}}|$ times a constant that we call `L1_penalty`.\n",
    "\n",
    "We will add both of these to the loss function:\n",
    "\\begin{equation}\n",
    "    L = (y - \\tilde{y})^2 + \\mathcal{R}_{L2} + \\mathcal{R}_{L1}\n",
    "\\end{equation}\n",
    "\n",
    "The parameters `L2_penalty` and `L1_penalty` are inputs to the train function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Bonus Coding Exercise 3.3: Add regularization to training \n",
    "\n",
    "We will create a new loss function that adds L1 and L2 regularization. \n",
    "In particular, you will:\n",
    "* add L2 loss penalty to the weights \n",
    "* add L1 loss penalty to the weights\n",
    "\n",
    "\n",
    "We will then train the network using this loss function. Full training will take a few minutes: if you want to train for just a few steps to speed up the code while iterating on your code, you can decrease the n_iter input from 500. \n",
    "\n",
    "Hint: since we are using `torch` instead of `np`, we will use `torch.abs` instead of `np.absolute`. You can use `torch.sum` or `.sum()` to sum over a tensor.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:29.005761Z",
     "iopub.status.busy": "2021-05-31T20:48:29.005158Z",
     "iopub.status.idle": "2021-05-31T20:48:29.038590Z",
     "shell.execute_reply": "2021-05-31T20:48:29.037843Z"
    }
   },
   "outputs": [],
   "source": [
    "def regularized_loss(output, target, weights, L2_penalty=0, L1_penalty=0):\n",
    "  \"\"\"loss function with L2 and L1 regularization\n",
    "\n",
    "  Args:\n",
    "    output (torch.Tensor): output of network\n",
    "    target (torch.Tensor): neural response network is trying to predict\n",
    "    weights (torch.Tensor): linear layer weights from neurons to hidden units (net.in_layer.weight)\n",
    "    L2_penalty : scaling factor of sum of squared weights\n",
    "    L1_penalty : scalaing factor for sum of absolute weights\n",
    "\n",
    "  Returns:\n",
    "    (torch.Tensor) mean-squared error with L1 and L2 penalties added\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  ##############################################################################\n",
    "  # TO DO: add L1 and L2 regularization to the loss function\n",
    "  raise NotImplementedError(\"Student exercise: complete regularized_loss\")\n",
    "  ##############################################################################\n",
    "\n",
    "  loss_fn = nn.NLLLoss()\n",
    "  loss = loss_fn(output, target)\n",
    "\n",
    "  L2 = L2_penalty * ...\n",
    "  L1 = L1_penalty * ...\n",
    "  loss += L1 + L2\n",
    "\n",
    "  return loss\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "n_classes = 20\n",
    "\n",
    "# Initialize network\n",
    "net = DeepNetSoftmax(n_neurons, 20, n_classes)  # use M=20 hidden units\n",
    "\n",
    "# Initialize built-in PyTorch MSE loss function\n",
    "loss_fn = regularized_loss\n",
    "\n",
    "# Here you can play with L2_penalty > 0, L1_penalty > 0\n",
    "train_loss, test_loss, predicted_test_labels = decode_orientation(net, n_classes, loss_fn,\n",
    "                                                                  resp_train, stimuli_train,\n",
    "                                                                  resp_test, stimuli_test,\n",
    "                                                                  n_iter=1000,\n",
    "                                                                  L2_penalty=1e-2,\n",
    "                                                                  L1_penalty=5e-4)\n",
    "\n",
    "# Plot results\n",
    "plot_decoded_results(train_loss, test_loss, stimuli_test, predicted_test_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:29.045441Z",
     "iopub.status.busy": "2021-05-31T20:48:29.044859Z",
     "iopub.status.idle": "2021-05-31T20:48:40.458788Z",
     "shell.execute_reply": "2021-05-31T20:48:40.458262Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "\n",
    "def regularized_loss(output, target, weights, L2_penalty=0, L1_penalty=0):\n",
    "  \"\"\"loss function with L2 and L1 regularization\n",
    "\n",
    "  Args:\n",
    "    output (torch.Tensor): output of network\n",
    "    target (torch.Tensor): neural response network is trying to predict\n",
    "    weights (torch.Tensor): linear layer weights from neurons to hidden units (net.in_layer.weight)\n",
    "    L2_penalty : scaling factor of sum of squared weights\n",
    "    L1_penalty : scalaing factor for sum of absolute weights\n",
    "\n",
    "  Returns:\n",
    "    (torch.Tensor) mean-squared error with L1 and L2 penalties added\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  loss_fn = nn.NLLLoss()\n",
    "  loss = loss_fn(output, target)\n",
    "\n",
    "  L2 = L2_penalty * torch.square(weights).sum()\n",
    "  L1 = L1_penalty * torch.abs(weights).sum()\n",
    "  loss += L1 + L2\n",
    "\n",
    "  return loss\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "\n",
    "n_classes = 20\n",
    "\n",
    "# Initialize network\n",
    "net = DeepNetSoftmax(n_neurons, 20, n_classes)  # use M=20 hidden units\n",
    "\n",
    "# Initialize built-in PyTorch MSE loss function\n",
    "loss_fn = regularized_loss\n",
    "\n",
    "# Here you can play with L2_penalty > 0, L1_penalty > 0\n",
    "train_loss, test_loss, predicted_test_labels = decode_orientation(net, n_classes, loss_fn,\n",
    "                                                                  resp_train, stimuli_train,\n",
    "                                                                  resp_test, stimuli_test,\n",
    "                                                                  n_iter=1000,\n",
    "                                                                  L2_penalty=1e-2,\n",
    "                                                                  L1_penalty=5e-4)\n",
    "\n",
    "# Plot results\n",
    "with plt.xkcd():\n",
    "  plot_decoded_results(train_loss, test_loss, stimuli_test, predicted_test_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It seems we were overfitting a little because we increased the accuracy a small amount by adding an L1 and L2 regularization penalty. What errors are still being made by the model?\n",
    "\n",
    "Let's see how the weights look after adding `L1_penalty > 0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:40.465350Z",
     "iopub.status.busy": "2021-05-31T20:48:40.463897Z",
     "iopub.status.idle": "2021-05-31T20:48:41.558193Z",
     "shell.execute_reply": "2021-05-31T20:48:41.557679Z"
    }
   },
   "outputs": [],
   "source": [
    "W_in = net.in_layer.weight.detach().numpy() # we can run detach and numpy to get a numpy array\n",
    "print('shape of W_in:')\n",
    "print(W_in.shape)\n",
    "\n",
    "plt.figure(figsize=(6,4))\n",
    "plt.imshow(W_in[:,isort], aspect='auto', cmap='bwr', vmin=-5e-3, vmax=5e-3)\n",
    "plt.xlabel('neurons')\n",
    "plt.ylabel('hidden units')\n",
    "plt.colorbar()\n",
    "plt.title('$W_{in}$')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The weights appear to be sparser than before."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "W2D1_Tutorial1",
   "provenance": [],
   "toc_visible": true
  },
  "kernel": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
